#团队协作
"""
角色：生成查询所有表的名字的sql代理tables_agent，生成查询单表的创表语句的代理info_agent，sql执行exec_agent、team团队配合
数据库代理----通过大模型生成sql语句，sql执行工具，查询数据库信息
1.初始planner/executor/gatherer
2.路由定义 router_edge ，当执行任务数量和收集信息数量相等时 转到 gatherer_node节点，否则装executor_node节点继续执行
3.画graph 节点和团队agent绑定
4.执行graph.invoke
"""
from exec_agent import Exec
from info_agent import Info
from tabel_agent import Table
from typing_extensions import TypedDict
from typing import Annotated, Optional
from operator import add
from model_utils import getLLM

_llm = getLLM()

_info = Info(_llm)
_table = Table(_llm)
_exec = Exec(_llm)

class DatebaseState(TypedDict):
    table_names:list[str]
    infos:Annotated[list[str],add]
    query:str
    table_name: str
    result:str


def _table_node(state):
    return _table({})

def _info_node(state):
    _table_name = state['table_name']
    _rt = _info({"table_name": _table_name})
    return {"infos":[_rt]}

def _exec_node(state):
    _infos = state["infos"]
    _query = state["query"]
    _rt = _exec({"query":_query,"infos":_infos})
    return {"result":_rt}

from langgraph.types import Send
from langgraph.graph import StateGraph,START,END
def _dispatch_table_names(state):
    sends = []
    for _table_name in state["table_names"]:
        sends.append(Send("_info_node", {"table_name": _table_name}))
    return sends

_builder = StateGraph(DatebaseState)

_builder.add_node("_table_node",_table_node)
_builder.add_node("_info_node",_info_node)
_builder.add_node("_exec_node",_exec_node)

_builder.add_edge(START,"_table_node")
# path_map没有参数，默认会连接之下的所有节点
_builder.add_conditional_edges("_table_node",_dispatch_table_names,
                               ['_info_node', '_exec_node', END])
_builder.add_edge("_info_node","_exec_node")
_builder.add_edge("_exec_node",END)

_graph = _builder.compile()
_graph.get_graph().draw_mermaid_png(output_file_path="team.png")

# _rt = _graph.invoke({"query":"我们公司男员工有多少"})
# _rt = _graph.invoke({"query":"我们公司女员工有多少"})
_rt = _graph.invoke({"query":"我们公司有多少员工"})
print(_rt.get("result"))

